package com.xzz.multipleMerge;

import java.util.*;

/**
 * @author: hhz
 * @create: 2022-01-14 11:51
 * 378. 有序矩阵中第 K 小的元素
 * 给你一个n x n矩阵matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
 * 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。
 **/
public class KthSmallest {
    public static void main(String[] args) {
        System.out.println(kthSmallest(new int[][]{{1, 5, 9}, {10, 11, 13}, {12, 13, 15}}, 8));
        System.out.println(kthSmallest(new int[][]{{-5}}, 1));
    }

    public static int kthSmallest(int[][] matrix, int k) {
        List<Integer> res = new ArrayList<>();
        PriorityQueue<int[]> queue = new PriorityQueue<>(Comparator.comparingInt(o -> matrix[o[0]][o[1]]));

        int len = matrix.length;
        for (int i = 0; i < len; i++) {
            queue.offer(new int[]{i, 0});
        }
        while (res.size() != k && !queue.isEmpty()) {
            int[] poll = queue.poll();
            res.add(matrix[poll[0]][poll[1]]);
            if (poll[1] + 1 < len) {
                queue.offer(new int[]{poll[0], poll[1] + 1});
            }
        }
        return res.get(k - 1);
    }
}
